import numpy as np
import elkai
import matplotlib.pyplot as plt


np.set_printoptions(suppress=True, linewidth=999)


def create_data_model(file_path):
    raw = np.loadtxt(file_path, skiprows=6, encoding='utf-8', dtype=float)
    res_data = dict()
    res_data['num_indexes'] = len(raw)
    res_data['x'] = raw[:, 1]
    res_data['y'] = raw[:, 2]
    i = np.arange(len(raw))
    j = i.reshape(len(i), -1)
    res_data['distance_matrix'] = np.sqrt((raw[:, 1][i] - raw[:, 1][j])**2 + (raw[:, 2][i] - raw[:, 2][j])**2)
    res_data['num_vehicles'] = 1
    res_data['depot'] = 0
    return res_data


if __name__ == '__main__':
    data = create_data_model('../../data/TSPLIB/att48.tsp.txt')
    route = elkai.solve_float_matrix(matrix=data['distance_matrix'], runs=1)
    route_length = sum(data['distance_matrix'][route[i - 1]][route[i]] for i in range(1, len(route))) + data['distance_matrix'][route[-1]][route[0]]
    plt.plot(data['x'][route+[route[0]]], data['y'][route+[route[0]]])
    plt.title(f'{route_length}')
    plt.show()


